using Test
using FLEXINVERT
using FLEXINVERT.PriorCovariance
using FLEXINVERT.Transformations
using FLEXINVERT.StateVector
using FLEXINVERT.CoreTypes
using FLEXINVERT.Settings
using FLEXINVERT.Regions
using FLEXINVERT.TransportInit
using LinearAlgebra

@testset "Covariance and Transformations Tests" begin

    # Create test configuration and domain
    config = Config(
        spec = "co2",
        datei = 20200301,
        datef = 20200302,
        w_edge_lon = 110.0,
        e_edge_lon = 112.0,
        s_edge_lat = 30.0,
        n_edge_lat = 32.0,
        xres = 0.5,
        yres = 0.5,
        nstep_nee = 2,
        nstep_nee_reg = 1,
        nstep_ff = 1,
        nstep_ff_reg = 1,
        nstep_ocn = 1,
        nstep_flx = 1,
        sigma_land = 1.0,
        sigma_ocean = 1.0,
        sigmatime = 1.0
    )

    domain = Domain(config)
    files = Files()
    FLEXINVERT.Regions.get_regions!(domain, files, config)
    date_info = TransportInit.initialize_dates(config)

    @testset "Spatial Correlation" begin
        spatial_corr = spatial_correlation(config, domain)

        @test size(spatial_corr) == (domain.nbox, domain.nbox)
        @test issymmetric(spatial_corr)

        # Diagonal should be 1.0
        @test all(diag(spatial_corr) .≈ 1.0)

        # Off-diagonal elements should be between 0 and 1
        @test all(0.0 .<= spatial_corr .<= 1.0)

        # Correlation should decrease with distance
        if domain.nbox >= 2
            # Find two boxes and check their correlation
            corr_12 = spatial_corr[1, 2]
            @test 0.0 <= corr_12 <= 1.0
        end

        @testset "Redundant mapping consistency" begin
            cfg = Config(
                sigma_land = 500_000.0,
                sigma_ocean = 1_500_000.0,
                w_edge_lon = 0.0,
                e_edge_lon = 4.0,
                s_edge_lat = 0.0,
                n_edge_lat = 3.0,
                xres = 1.0,
                yres = 1.0
            )

            dom_redundant = Domain(; rllx=0.0, rlly=0.0, rdx=1.0, rdy=1.0,
                                     nxregrid=4, nyregrid=3, nbox=3,
                                     nbox_xy=zeros(Int, 4, 3),
                                     lsm_box=ones(Int, 4, 3))
            dom_redundant.nbox_xy[:, 1] .= 1
            dom_redundant.nbox_xy[:, 2] .= Int[1, 1, 3, 3]
            dom_redundant.nbox_xy[:, 3] .= 2
            dom_redundant.lsm_box[:, 3] .= 0
            dom_redundant.nbox = maximum(dom_redundant.nbox_xy)

            expected = zeros(Float64, 3, 3)
            stats = PriorCovariance.compute_box_summaries(dom_redundant)
            cutoff = 1.0e-4
            for i in 1:3
                expected[i, i] = 1.0
                stats_i = stats[i]
                for j in (i + 1):3
                    stats_j = stats[j]
                    cor = 0.0
                    if stats_i.land_fraction > 0 && stats_j.land_fraction > 0
                        dist = PriorCovariance.haversine_distance(stats_i.land_lat, stats_i.land_lon,
                                                                  stats_j.land_lat, stats_j.land_lon)
                        weight = stats_i.land_fraction * stats_j.land_fraction
                        cor += weight * exp(-dist / cfg.sigma_land)
                    end
                    if stats_i.ocean_fraction > 0 && stats_j.ocean_fraction > 0
                        dist = PriorCovariance.haversine_distance(stats_i.ocean_lat, stats_i.ocean_lon,
                                                                  stats_j.ocean_lat, stats_j.ocean_lon)
                        weight = stats_i.ocean_fraction * stats_j.ocean_fraction
                        cor += weight * exp(-dist / cfg.sigma_ocean)
                    end
                    if cor < cutoff
                        cor = 0.0
                    end
                    expected[i, j] = expected[j, i] = cor
                end
            end

            corr_redundant = spatial_correlation(cfg, dom_redundant)
            @test corr_redundant ≈ expected atol=1e-12
        end
    end

    @testset "Temporal Correlation" begin
        # Single time step
        temporal_corr_1 = temporal_correlation(config, 1)
        @test size(temporal_corr_1) == (1, 1)
        @test temporal_corr_1[1, 1] ≈ 1.0

        # Multiple time steps
        temporal_corr_3 = temporal_correlation(config, 3)
        @test size(temporal_corr_3) == (3, 3)
        @test issymmetric(temporal_corr_3)

        # Diagonal should be 1.0
        @test all(diag(temporal_corr_3) .≈ 1.0)

        # Lag-1 correlation should be less than 1
        @test temporal_corr_3[1, 2] < 1.0
        @test temporal_corr_3[1, 2] > 0.0
    end

    @testset "Covariance Matrix Construction" begin
        cov_matrix = build_prior_covariance(config, domain)

        @test cov_matrix isa CovarianceMatrix
        @test cov_matrix.nbox == domain.nbox
        @test cov_matrix.n_modes > 0
        @test cov_matrix.n_modes <= domain.nbox

        # Eigenvalues should be positive
        @test all(cov_matrix.eigenvalues .> 0)

        # Eigenvalues should be sorted in descending order
        @test issorted(cov_matrix.eigenvalues, rev=true)

        # Explained variance should be reasonable
        @test 0.0 < cov_matrix.explained_variance <= 1.0

        # Eigenvectors should be orthonormal
        U = cov_matrix.eigenvectors
        @test norm(U' * U - I) < 1e-10
    end

    @testset "Phi-Chi Transformations" begin
        cov_matrix = build_prior_covariance(config, domain)

        # Test round-trip in control space
        phi_test = randn(cov_matrix.n_modes)
        chi_intermediate = phi2chi(phi_test, cov_matrix)
        phi_recovered = chi2phi(chi_intermediate, cov_matrix)

        @test phi_test ≈ phi_recovered rtol=1e-10

        # Test dimensions
        @test length(chi_intermediate) == domain.nbox
        @test length(phi_recovered) == cov_matrix.n_modes

        # Test in-place operations
        chi_inplace = zeros(domain.nbox)
        phi2chi!(chi_inplace, phi_test, cov_matrix)
        @test chi_inplace ≈ chi_intermediate

        phi_inplace = zeros(cov_matrix.n_modes)
        chi2phi!(phi_inplace, chi_intermediate, cov_matrix)
        @test phi_inplace ≈ phi_test rtol=1e-10
    end

    @testset "Transform Validation" begin
        cov_matrix = build_prior_covariance(config, domain)

        # Validate transforms
        @test validate_transforms(cov_matrix)

        # Test effective rank
        rank = effective_rank(cov_matrix)
        @test rank == cov_matrix.n_modes

        # Test condition number
        cond_num = condition_number(cov_matrix)
        @test cond_num > 1.0
        @test isfinite(cond_num)
    end

    @testset "State Vector Operations" begin
        cov_matrix = build_prior_covariance(config, domain)
        state = initialize_state(config, domain, cov_matrix)

        @test state isa InversionState
        @test state.n_control == cov_matrix.n_modes
        @test state.n_physical == domain.nbox
        @test state.iteration == 0

        # Test state consistency
        @test validate_state_consistency(state)

        # Test state update
        dphi = 0.1 * randn(state.n_control)
        update_state!(state, dphi)

        @test state.phi ≈ dphi
        @test validate_state_consistency(state)

        # Test state statistics
        stats = get_state_statistics(state)
        @test stats.iteration == 0
        @test stats.n_control == cov_matrix.n_modes
        @test stats.n_physical == domain.nbox

        # Test optimization history
        add_to_optimization_history!(state, 100.0, 10.0)
        @test state.iteration == 1
        @test length(state.cost_history) == 1
        @test state.cost_history[1] == 100.0

        # Test reset to prior
        reset_to_prior!(state)
        @test state.iteration == 0
        @test length(state.cost_history) == 0
        @test all(state.phi .== 0.0)
    end

    @testset "Distance Calculations" begin
        # Test haversine distance
        # Distance between (0,0) and (0,1) should be approximately 111 km
        dist = PriorCovariance.haversine_distance(0.0, 0.0, 1.0, 0.0)
        @test dist ≈ 1.11195e5 rtol=0.01

        # Distance from point to itself should be zero
        dist_zero = PriorCovariance.haversine_distance(30.0, 110.0, 30.0, 110.0)
        @test dist_zero ≈ 0.0 atol=1e-10
    end

    @testset "Error Variance Calculations" begin
        variances = PriorCovariance.get_error_variances(config, domain, nothing, date_info)

        @test length(variances) == domain.nbox
        @test all(variances .> 0.0)

        # Variances should scale with area
        expected = (0.5 * domain.area_box[1])^2
        @test variances[1] ≈ expected
    end
end
